import matplotlib.pyplot as plt
import numpy as np
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression


"""构造数据"""
m = 100
X = 6*np.random.rand(m, 1) - 3
y = 0.5*X**2+X+np.random.randn(m, 1)  # 高斯抖动
"""数据作图"""
plt.plot(X, y, 'b.')
plt.xlabel("x_1")
plt.ylabel("y")
plt.axis([-3, 3, -5, 15])
plt.show()

"""初始化"""
poly_features = PolynomialFeatures(degree=2, include_bias=False)  # 只对x做偏置
X_poly = poly_features.fit_transform(X)

# 线性回归
lin_reg = LinearRegression()
lin_reg.fit(X_poly, y)
print(lin_reg.coef_)
print(lin_reg.intercept_)

"""构造测试集合"""
X_new = np.linspace(-3, 3, 100).reshape(100, 1)
X_new_poly = poly_features.transform(X_new)
y_new = lin_reg.predict(X_new_poly)

"""作图看看"""
plt.plot(X, y, 'b.')
plt.plot(X_new, y_new, 'r--', label='prediction')
plt.axis([-3, 3, -5, 10])
plt.legend()
plt.show()
